#!/bin/bash

# 检查参数数量
if [ $# -eq 0 ]; then
    echo "Usage: $0 [cpt|sft <model_name>|chat <model_name>|export <model_name>]"
    echo ""
    echo "Available commands:"
    echo "  cpt                        - Run pretrain training"
    echo "  sft <model_name>     - Run LoRA SFT training for specified model"
    echo "  export <model_name>  - Export LoRA model for specified model"
    echo "  chat <model_name>         - Start chat interface with specified model"
    echo ""
    echo "Examples:"
    echo "  $0 sft sft_phase_1"
    echo "  $0 chat sft_phase_1"
    echo "  $0 export sft_phase_1"
    exit 1
fi

cd "LLaMA-Factory"

# 根据参数选择执行的命令
case "$1" in
    "cpt")
        echo "Starting pretrain training..."
        llamafactory-cli train ./configs/cpt.yaml
        ;;
    "sft")
        if [ -z "$2" ]; then
            echo "Error: sft command requires a model name"
            echo "Usage: $0 sft <model_name>"
            echo "Example: $0 sft phase_1"
            exit 1
        fi
        echo "Starting LoRA SFT training for model: $2"
        # 这里可以根据需要修改配置文件路径或传递参数
        llamafactory-cli train ./configs/$2_sft.yaml
        ;;
    "export")
        if [ -z "$2" ]; then
            echo "Error: export command requires a model name"
            echo "Usage: $0 export <model_name>"
            echo "Example: $0 export phase_1"
            exit 1
        fi
        echo "Exporting LoRA model for: $2"
        # 这里可以根据需要修改配置文件路径或传递参数
        llamafactory-cli export ./configs/$2_merge.yaml
        ;;
    "chat")
        if [ -z "$2" ]; then
            echo "Error: chat command requires a model name"
            echo "Usage: $0 chat <model_name>"
            echo "Example: $0 chat sft_phase_1"
            exit 1
        fi
        echo "Starting chat interface with model: $2"
        # 这里可以根据需要修改配置文件路径或传递参数
        llamafactory-cli export ./configs/$2_chat.yaml
        ;;

    "eval")
        if [ -z "$2" ]; then
            echo "Error: eval command requires a model name"
            echo "Usage: $0 eval <model_name>"
            echo "Example: $0 eval phase_1"
            exit 1
        fi
        echo "Starting evaluation for model: $2"
        llamafactory-cli eval ./configs/$2_eval.yaml
        ;;
    *)
        echo "Error: Unknown command '$1'"
        echo "Available commands: cpt, sft <model_name>, chat <model_name>, export <model_name>"
        echo "Examples:"
        echo "  $0 sft phase_1"
        echo "  $0 export phase_1"
        echo "  $0 chat sft_phase_1"
        exit 1
        ;;
esac